from client import es_client


async def get_iot_run_data(depart_id: int, start_time, end_time, imei_list: list = None):
    # TODO: Implement the function to fetch IoT run data from Elasticsearch
    query = {
        "size": 0,
        "query": {
            "bool": {
                "must": [
                    {
                        "range": {
                            "startTime": {
                                "gte": start_time,
                                "lte": end_time
                            }
                        }
                    }
                    # ,{
                    #     "term": {
                    #         "dataAuthority.depart": {
                    #             "value": depart_id
                    #         }
                    #     }
                    # }
                ],

            }
        },
        "aggs": {
            "by_imei": {
                "terms": {
                    "field": "imei",
                    "size": 10000
                },
                "aggs": {
                    "by_day": {
                        "date_histogram": {
                            "field": "startTime",
                            "calendar_interval": "day"
                        },
                        "aggs": {
                            "total_runtime_hours": {
                                "sum": {
                                    "field": "intervals"
                                }
                            },
                            "runtime_in_hours": {
                                "bucket_script": {
                                    "buckets_path": {
                                        "total_runtime": "total_runtime_hours"
                                    },
                                    "script": "params.total_runtime  / (1000 * 60 * 60)"
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    if imei_list is not None:
        query['query']['bool']['must'].append({"terms": {"imei": imei_list}})

    es = await es_client.get_es_client()
    result = await es.search(index="i_iot_run", query=query)
    result = result['aggregations']['by_imei']['buckets']
    run_data = []
    for bucket in result:
        print(bucket)
        for day_bucket in bucket['by_day']['buckets']:
            print(day_bucket)
            if int(day_bucket.get('doc_count')) > 0:
                run_data.append({
                    'imei': bucket['key'],
                    'date': day_bucket['key_as_string'],
                    'runtime_hours': float(day_bucket['runtime_in_hours']['value']).__round__(2)

                })
    return run_data


async def get_iot_run_data_detail(term: dict, start_time, end_time):
    # TODO: Implement the function to fetch IoT run data from Elasticsearch
    query = {
        "size": 1000,
        "_source": ['imei', 'categoryName', 'startTime', 'endTime', 'intervals', 'createDate'],
        "query": {
            "bool": {
                "must": [
                    {
                        "term": term
                    },
                    {
                        "range": {
                            "startTime": {
                                "gte": start_time,
                                "lte": end_time
                            }
                        }
                    }
                ]
            }
        }
    }
    es = await es_client.get_es_client()
    return await es.search_documents(index="i_iot_run", query=query)


async def get_equipment_list(query_bool: dict):
    query = {
        "size": 10000,
        "_source": ["imei", "member.name", "project.name", "city.name", "workAmount", "region.name", "workUnit",
                    "memberTypeName", "equipmentAlias", "createDate"],
        "query": {
            "bool": query_bool
        }
    }
    es = await es_client.get_es_client()
    result = await es.search(index="b_equipment_entity", query=query)
    equipment_info_list = []
    for item in result['hits']['hits']:
        print(item)
        equipment_info_list.append({
            'imei': item['_source']['imei'],
            'member_name': item['_source']['member']['name'],
            'region_name': item['_source']['region']['name'],
            'project_name': item['_source']['project']['name'],
            'city_name': item['_source']['city']['name'],
            'work_amount': item['_source']['workAmount'],
            'work_unit': item['_source']['workUnit'],
            'type': item['_source']['memberTypeName'],
            'createDate': item['_source']['createDate']
            # 'alias': item['_source']['equipmentAlias']
        })

    return equipment_info_list
